# -*- coding: utf-8 -*-
# @Time    : 2019/10/25 10:27
# @Author  : CHEN Wang
# @Site    : 
# @File    : fund_classification.py
# @Software: PyCharm

"""
脚本说明: 获取基金的分类相关信息
"""
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.db_operator import db_conn, my_mysql
from quant_researcher.quant.project_tool import hammer
from quant_researcher.quant.project_tool.celebrity import get_lst_condition

F_C = 'fund_code'
C_R = 'category_rank'
C_N = 'category_name'
C_C = 'category_code'
T_CLASSIFICATION_INFO = 'mf_bd_fndclsinfo'
T_CLASSIFICATION_NAME = 'mf_bd_fndclsstd'
T_CLASSIFICATION_INFO_WIND = 'mf_bd_fundclsinfo_wind'


def get_codes_by_category_rank_and_code(cls_rank=None, cls_code=None, conn=None, **kwargs):
    """
    根据基金分类层级和分类代码查询对应的所有基金，支持不同家的分类

    :param cls_rank: list 或 int 或 None，分类层级
        list 查询多个层级的；int 查询单个层级；None，SQL不加分类层级限定
    :param cls_code: list 或 str 或 None，分类代码
        list 查询多个代码的；str 查询单个代码；None，SQL不加分类代码限定
    :param conn: 连接对象，数据库连接
    :param kwargs: 关键字参数
        select，list 或 str，具体参照查询函数 my_mysql.read_v2，默认：None
        clsstd_code, 采用哪家的分类，默认为'thinkive', 还支持'wind', 'haitong', 'shanghai'
        end_date, 获取哪天的分类
    :return: pd.DataFrame
    """
    tmp_conn = db_conn.get_basic_data_conn() if conn is None else conn
    select = kwargs.pop('select', None)
    if select is None:
        if cls_rank is None and cls_code is None:
            select = [f'distinct {F_C}']
        else:
            select = [F_C]

    end_date = kwargs.pop('end_date', None)
    clsstd_code = kwargs.pop('clsstd_code', 'thinkive')
    if clsstd_code != 'thinkive':  # 如果指定了采用其他分类体系
        if end_date is None:
            end_date = my_mysql.read_v2(select='max(end_date)',
                                        sfrom=T_CLASSIFICATION_INFO_WIND, conn=tmp_conn).values[0][0]
        df = my_mysql.read_v2(select=['fund_code', 'fund_sname', 'end_date'],
                              sfrom=T_CLASSIFICATION_INFO_WIND,
                              where=[f"fund_type_2 = '{cls_code}'",
                                     f"clsstd_code = '{clsstd_code}'",
                                     f"end_date = '{end_date}'"],
                              conn=tmp_conn)
    else:  # 如果是采用思迪的分类体系
        # where = [f"clsstd_code='{clsstd_code}' and state=0 and classify_date='{end_date}'"] #
        # 后续需要把历史的分类信息补充完整，就可以使用上面这个
        where = [f"clsstd_code='{clsstd_code}' and state=0"]
        if end_date is None:
            end_date = my_mysql.read_v2(select='max(classify_date)',
                                        sfrom=T_CLASSIFICATION_INFO, conn=tmp_conn).values[0][0]
        if cls_code is None:
            condition_code = None
        else:
            if not isinstance(cls_code, list):
                cls_code = [cls_code]
            condition_code = ' or '.join([f'{C_C}={x}' for x in cls_code])
        if cls_rank is None:
            condition_rank = None
        else:
            if not isinstance(cls_rank, list):
                cls_rank = [cls_rank]
            condition_rank = ' or '.join([f'{C_R}={x}' for x in cls_rank])
        if condition_code is not None:
            where.append(condition_code)
        if condition_rank is not None:
            where.append(condition_rank)
        if where:
            where = ' and '.join([f'({x})' for x in where])
        else:
            where = None
        df = my_mysql.read_v2(select=select, sfrom=T_CLASSIFICATION_INFO, where=where, conn=tmp_conn)
    hammer.slim(df, name=f'get_codes_by_category_rank({cls_rank})_and_code({cls_code})', a_f=1)
    if conn is None:
        tmp_conn.close()

    return df


def get_distinct_cls_codes_by_category_rank(category_rank, conn=None):
    """
    根据不同的分类层级，查询该层级下的所有 不同的（distinct） 分类代码

    :param category_rank，str，分类层级
    :param conn，数据库连接对象，默认：None（使用默认配置的数据库连接）
    :return pd.DataFrame
    """
    conn = db_conn.get_global_basic_data_conn() if conn is None else conn
    select = f'distinct {C_C}'
    sfrom = T_CLASSIFICATION_INFO
    where = f'{C_R}={category_rank}'
    df = my_mysql.read_v2(select=select, sfrom=sfrom, where=where, conn=conn)
    hammer.slim(df, name=f'{select}，{where}', a_f=1)
    return df


def get_fund_cls_info(fund_code=None, **kwargs):
    """
    查询思迪基金分类中，指定基金的一级、二级分类信息

    :param fund_code: 基金代码，支持str、list、None
    :param kwargs:
        - category_rank，str，分类层级，默认：None，不加限定
        - other_filter，str，传入SQL的其他筛选条件，默认：None，不加限定
        - conn，数据库连接对象，默认：None（使用默认配置的数据库连接）
        - select，list 或 str，具体参照查询函数 my_mysql.read_v2
            默认：基金代码，分类代码
    :return pd.DataFrame
    """
    category_rank = kwargs.pop('category_rank', None)  # 对应的是几级分类
    other_filter = kwargs.pop('other_filter', None)  # 其他可能的筛选条件
    conn = kwargs.pop('conn', None)
    select = kwargs.pop('select', None)
    tmp_conn = db_conn.get_basic_data_conn() if conn is None else conn

    select = [F_C, C_C] if select is None else select
    where = [f'clsstd_code=\'thinkive\' and state=0']
    if isinstance(fund_code, str):
        where.append(f'{F_C}=\'{fund_code}\'')
    elif isinstance(fund_code, list):
        where.append(get_lst_condition(fund_code, F_C))
    elif fund_code is None:
        pass
    else:
        LOG.error('目前fund_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError
    if category_rank is not None:
        where.append(f'{C_R}={category_rank}')
    if other_filter is not None:
        where.append(other_filter)
    sfrom = T_CLASSIFICATION_INFO

    df = my_mysql.read_v2(sfrom=sfrom, select=select, where=where, conn=tmp_conn)
    row_count = df.shape[0]
    # how_many_classes_4_a_code = 2
    if fund_code is not None:
        if row_count == 0:
            LOG.warning(f'没有在 {sfrom} 中找到 {fund_code} 的信息，查询条件：{where}')
            return
        # if category_rank is not None and row_count > 1:
        #     err_msg = f'传入了基金代码={fund_code}，对应 {category_rank} 级分类，' \
        #               f'却找到的多余 1 条记录，查询条件：{where}'
        #     LOG.error(err_msg)
        #     raise RuntimeError(err_msg)
        # if category_rank is None \
        #         and (row_count == 1 or row_count > how_many_classes_4_a_code):
        #     err_msg = f'目前应该有 {how_many_classes_4_a_code} 级分类，' \
        #               f'但是 {fund_code} 却只找到了 {row_count} 条记录，' \
        #               f'查询条件：{where}'
        #     LOG.error(err_msg)
        #     raise RuntimeError(err_msg)
    # hammer.slim(df, a_f=1)
    if conn is None:
        tmp_conn.close()

    return df


def get_all_fund_cls_info(**kwargs):
    """
    得到每个基金的一级、二级分类信息，以及没有二级分类的基金用一级分类填充

    :return: pd.DataFrame
                column: type_1-一级分类
                column: type_2-二级分类
                column: type-没有二级分类的基金用一级分类填充
    """
    conn = kwargs.pop('conn', None)
    tmp_conn = db_conn.get_basic_data_conn() if conn is None else conn

    cls_1_df = get_fund_cls_info(category_rank=1, select=['fund_code', 'category_code as `type_1`'], conn=tmp_conn)
    cls_2_df = get_fund_cls_info(category_rank=2, select=['fund_code', 'category_code as `type_2`'], conn=tmp_conn)
    cls_df = cls_1_df.merge(cls_2_df, how='left', on='fund_code')
    cls_df['type'] = cls_df['type_2'].fillna(cls_df['type_1'])
    if conn is None:
        tmp_conn.close()

    return cls_df


def get_cls_name(category_rank=None):
    """
    得到基金分类代码对应的名称

    :param category_rank: 类型级别，可选1-一级分类，2-二级分类，None-一级、二级同时输出
    :return: pd.DataFrame
    """
    conn = db_conn.get_basic_data_conn()
    df = my_mysql.read_v2(select=['category_code', 'category_name', 'category_rank'],
                          sfrom=T_CLASSIFICATION_NAME, conn=conn)
    if category_rank is not None:
        df = df[df['category_rank'] == category_rank]

    df = df[['category_code', 'category_name']]
    conn.close()
    return df


if __name__ == '__main__':
    # get_fund_cls_info(category_rank=2, select=['fund_code', 'category_code as `二级分类`'])
    get_fund_cls_info(fund_code=['110022', '003095'], category_rank=1, select=['fund_code', 'category_code as `一级分类`'])
    # conn = db_conn.get_derivative_data_conn()
    # get_codes_by_category_rank_and_code(cls_rank=None, cls_code='增强指数型基金', clsstd_code='wind',
    #                                     conn=conn, end_date=None)
